Feature/hiera #418

benjijamorris · 2024-08-19T21:32:38Z

What does this PR do?

add 2d and 3d Hiera Models
refactor mae/hiera to encoder/decoder.py instead of model-specific files
unify patchify code
add tests
updates to JEPA to accommodate refactoring

Before submitting

Did you make sure title is self-explanatory and the description concisely explains the PR?
Did you make sure your PR does only one thing, instead of bundling different changes together?
Did you list all the breaking changes introduced by this pull request?
Did you test your PR locally with pytest command?
Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃

…nto main

ritvikvasan

This is a lot!! Loved reading it...As always, minor comments

ritvikvasan · 2024-08-21T16:20:32Z

configs/model/im2im/hiera.yaml

+backbone:
+  _target_: cyto_dl.nn.vits.mae.HieraMAE
+  spatial_dims: ${spatial_dims}
+  patch_size: 2 # patch_size* num_patches should be your patch shape


should be image shape?

can this be a list for ZYX patch size?

yes - the terminology is confusing here haha. "patch" = the small crop extracted from your original image, but "patch" is also the tokenized component of the image fed into the network. The patch size can be either an int (repeated for each spatial dim) or a list of size spatial_dims

ritvikvasan · 2024-08-21T16:21:37Z

configs/model/im2im/hiera.yaml

+  spatial_dims: ${spatial_dims}
+  patch_size: 2 # patch_size* num_patches should be your patch shape
+  num_patches: 8 # patch_size * num_patches = img_shape
+  num_mask_units: 4 #img_shape / num_mask_units = size of each mask unit in pixels, num_patches/num_mask_units = number of patches permask unit


Clarify what a mask unit is here?

ritvikvasan · 2024-08-21T17:39:06Z

configs/model/im2im/hiera.yaml

+  architecture:
+    # mask_unit_attention blocks - attention is only done within a mask unit and not across mask units
+    # the total amount of q_stride across the architecture must be less than the number of patches per mask unit
+    - repeat: 1


what is repeat?

ritvikvasan · 2024-08-21T17:40:07Z

configs/model/im2im/hiera.yaml

+    # self attention transformer - attention is done across all patches, irrespective of which mask unit they're in
+    - repeat: 2
+      num_heads: 4
+      self_attention: True


so last layer is global attention and first 2 layers are local attention? Is 3 layers the recommended hierarchy?

correct. 3 layers is small enough to test quickly. All of the models with unit tests are tiny by default in the configs and I have somewhere in the docs that you should increase the model size if you want good performance.

ritvikvasan · 2024-08-21T17:49:07Z

cyto_dl/nn/vits/blocks/masked_unit_attention.py

+            if self.spatial_dims == 3:
+                q = reduce(
+                    q,
+                    "b n h (n_patches_z q_stride_z n_patches_y q_stride_y n_patches_x q_stride_x) c ->b n h (n_patches_z n_patches_y n_patches_x) c",


can you use the same nomenclature here? e.g. n = num_mask_units = mask_units, num_heads = h = n_heads

c = head_dim

ritvikvasan · 2024-08-21T17:54:27Z

cyto_dl/nn/vits/blocks/masked_unit_attention.py

+        self.spatial_dims = spatial_dims
+        self.num_heads = num_heads
+        self.head_dim = dim_out // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5


this isn't used anywhere

ritvikvasan · 2024-08-21T18:03:20Z

cyto_dl/nn/vits/blocks/masked_unit_attention.py

+        # change dimension and subsample within mask unit for skip connection
+        x = self.proj(x_norm)
+
+        x = x + self.drop_path(self.attn(x_norm))


Does dim_out = dim for skip connection with attention?

Good question - each block specified in the architecture argument doubles the embedding dimension and halves the size of the mask unit. This doubling/pooling happens on the last repeat of the block, so dim_out=dim for all repeats except the last. I updated the docstring with an example.

ritvikvasan · 2024-08-21T18:05:55Z

cyto_dl/nn/vits/blocks/masked_unit_attention.py

+        dim_out: int,
+        heads: int,
+        spatial_dims: int = 3,
+        mlp_ratio: float = 4.0,


what is mlp_ratio? add to docstring?

ritvikvasan · 2024-08-21T18:14:18Z

cyto_dl/nn/vits/blocks/patchify/patchify_hiera.py

+
+
+class PatchifyHiera(PatchifyBase):
+    """Class for converting images to a masked sequence of patches with positional embeddings."""


to "mask units" instead of masked sequence? since that's what a regular patchify does?

ritvikvasan · 2024-08-21T18:30:51Z

cyto_dl/nn/vits/utils.py

@@ -40,3 +47,8 @@ def get_positional_embedding(
            cls_token = torch.zeros(1, 1, emb_dim)
            pe = torch.cat([cls_token, pe], dim=0)
        return torch.nn.Parameter(pe, requires_grad=False)
+
+
+def validate_spatial_dims(spatial_dims, tuples):


I feel like the code might be clearer by not having this be a separate function and just calling these 2 lines in every class? I thought this function was doing a lot more based on the name (like some math to check that the spatial dimensions of each patch and mask is correct). What do you think?

I'd prefer renaming it to something clearer rather than repeating the code, maybe match_tuple_to_spatial_dims?

sounds good

ritvikvasan · 2024-08-21T18:37:02Z

windows tests always seem to be failing. any ideas why?

benjijamorris · 2024-08-21T20:13:17Z

the windows tests are just way slower for some reason... all the tests pass but I set a 70 minute time out so we don't rack up crazy costs.

Benjamin Morris and others added 24 commits May 2, 2024 10:02

Bump version: 0.1.5 → 0.1.6

6ed231d

Merge branch 'main' of https://github.com/AllenCellModeling/cyto-dl i…

004b823

…nto main

Merge branch 'main' of https://github.com/AllenCellModeling/cyto-dl i…

0c607d9

…nto main

add hiera

2c0c275

start of mask2former

50646ae

first take at transfomer

22ddadc

fix dimensionality, now updating instance queries instead of mask

693c514

give instance queries own dim

913416a

add mask creation

a561aa4

Merge branch 'main' into feature/hiera

e5591d4

remove experimental code

9f6591c

update to base patchify

a773f3b

merge main

70a0437

wip

23f4b94

update configs

357166f

update patchify

a12d8cd

rearrange encoder/decoder/mae

f652138

add 2d hiera

7ee08e3

2d masked unit attention

71b8096

precommit

038d573

update configs

5098732

update hiera model config

453f0e5

update deafults

ec07e39

update tests

128baa5

benjijamorris requested a review from ritvikvasan August 19, 2024 21:32

Benjamin Morris added 5 commits August 19, 2024 14:33

delete patchify_conv

3c97933

fix jepa tests

1d4a76d

update mask transform

663e52c

add predictor

30a0526

precommit

e6b58df

ritvikvasan previously approved these changes Aug 21, 2024

View reviewed changes

update with ritviks comments

84717b7

benjijamorris dismissed ritvikvasan’s stale review via 84717b7 August 22, 2024 22:51

benjijamorris requested a review from ritvikvasan August 22, 2024 23:40

ritvikvasan previously approved these changes Aug 23, 2024

View reviewed changes

replace all function names

316a93a

benjijamorris dismissed ritvikvasan’s stale review via 316a93a August 23, 2024 18:58

benjijamorris requested a review from ritvikvasan August 23, 2024 22:22

ritvikvasan approved these changes Aug 23, 2024

View reviewed changes

benjijamorris merged commit 48e7cb9 into main Aug 26, 2024
4 of 6 checks passed

benjijamorris deleted the feature/hiera branch August 26, 2024 17:42

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Feature/hiera #418

Feature/hiera #418

benjijamorris commented Aug 19, 2024 •

edited

Loading

ritvikvasan left a comment

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

benjijamorris Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

benjijamorris Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

benjijamorris Aug 22, 2024

ritvikvasan Aug 22, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

ritvikvasan Aug 21, 2024

benjijamorris Aug 22, 2024

ritvikvasan Aug 22, 2024

ritvikvasan commented Aug 21, 2024

benjijamorris commented Aug 21, 2024



		class PatchifyHiera(PatchifyBase):
		"""Class for converting images to a masked sequence of patches with positional embeddings."""

Feature/hiera #418

Feature/hiera #418

Conversation

benjijamorris commented Aug 19, 2024 • edited Loading

What does this PR do?

Before submitting

Did you have fun?

ritvikvasan left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ritvikvasan commented Aug 21, 2024

benjijamorris commented Aug 21, 2024

benjijamorris commented Aug 19, 2024 •

edited

Loading